{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   age          workclass  fnlwgt   education  education_num  \\\n",
      "0   39          State-gov   77516   Bachelors             13   \n",
      "1   50   Self-emp-not-inc   83311   Bachelors             13   \n",
      "2   38            Private  215646     HS-grad              9   \n",
      "3   53            Private  234721        11th              7   \n",
      "4   28            Private  338409   Bachelors             13   \n",
      "\n",
      "        marital_status          occupation    relationship    race      sex  \\\n",
      "0        Never-married        Adm-clerical   Not-in-family   White     Male   \n",
      "1   Married-civ-spouse     Exec-managerial         Husband   White     Male   \n",
      "2             Divorced   Handlers-cleaners   Not-in-family   White     Male   \n",
      "3   Married-civ-spouse   Handlers-cleaners         Husband   Black     Male   \n",
      "4   Married-civ-spouse      Prof-specialty            Wife   Black   Female   \n",
      "\n",
      "   capital_gain  capital_loss  hours_per_week  native_country high_income  \n",
      "0          2174             0              40   United-States       <=50K  \n",
      "1             0             0              13   United-States       <=50K  \n",
      "2             0             0              40   United-States       <=50K  \n",
      "3             0             0              40   United-States       <=50K  \n",
      "4             0             0              40            Cuba       <=50K  \n"
     ]
    }
   ],
   "source": [
    "import pandas\n",
    "\n",
    "# Set index_col to False to avoid pandas thinking that the first column is row indexes (it's age).\n",
    "columns = [\"age\", \"workclass\", \"fnlwgt\", \"education\", \"education_num\", \"marital_status\", \"occupation\", \"relationship\", \"race\", \"sex\", \n",
    "           \"capital_gain\", \"capital_loss\", \"hours_per_week\", \"native_country\", \"high_income\"]\n",
    "#income = pandas.read_csv(\"income.csv\", index_col=False,names=columns)\n",
    "income = pandas.read_csv(\"income.csv\", names=columns)\n",
    "print(income.head(5))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[' State-gov' ' Self-emp-not-inc' ' Private' ' Federal-gov' ' Local-gov'\n",
      " ' ?' ' Self-emp-inc']\n",
      "[State-gov, Self-emp-not-inc, Private, Private, Private, ..., Private, Private, Private, Self-emp-not-inc, Private]\n",
      "Length: 629\n",
      "Categories (7, object): [?, Federal-gov, Local-gov, Private, Self-emp-inc, Self-emp-not-inc, State-gov]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "e:\\python27\\lib\\site-packages\\ipykernel\\__main__.py:3: FutureWarning: Categorical.from_array is deprecated, use Categorical instead\n",
      "  app.launch_new_instance()\n",
      "e:\\python27\\lib\\site-packages\\ipykernel\\__main__.py:8: FutureWarning: Categorical.from_array is deprecated, use Categorical instead\n"
     ]
    }
   ],
   "source": [
    "# Convert a single column from text categories into numbers.\n",
    "print income[\"workclass\"].unique()\n",
    "col = pandas.Categorical.from_array(income[\"workclass\"])\n",
    "print col\n",
    "income[\"workclass\"] = col.codes\n",
    "#print(income[\"workclass\"].head(5))\n",
    "for name in [\"education\", \"marital_status\", \"occupation\", \"relationship\", \"race\", \"sex\", \"native_country\", \"high_income\"]:\n",
    "    col = pandas.Categorical.from_array(income[name])\n",
    "    income[name] = col.codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(19, 15)\n",
      "(610, 15)\n"
     ]
    }
   ],
   "source": [
    "# Enter your code here.\n",
    "private_incomes = income[income[\"workclass\"] == 4]\n",
    "public_incomes = income[income[\"workclass\"] != 4]\n",
    "\n",
    "print(private_incomes.shape)\n",
    "print(public_incomes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.528771237955\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "# We'll do the same calculation we did above, but in Python.\n",
    "# Passing 2 as the second parameter to math.log will take a base 2 log.\n",
    "entropy = -(2/5.0 * math.log(2/5.0, 2) + 3/5 * math.log(3/5.0, 2))\n",
    "print(entropy)\n",
    "prob_0 = income[income[\"high_income\"] == 0].shape[0] / float(income.shape[0])\n",
    "prob_1 = income[income[\"high_income\"] == 1].shape[0] / float(income.shape[0])\n",
    "income_entropy = -(prob_0 * math.log(float(prob_0), 2) + prob_1 * math.log(float(prob_1), 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.756141116271\n"
     ]
    }
   ],
   "source": [
    "import numpy\n",
    "\n",
    "def calc_entropy(column):\n",
    "    \"\"\"\n",
    "    Calculate entropy given a pandas Series, list, or numpy array.\n",
    "    \"\"\"\n",
    "    # Compute the counts of each unique value in the column.\n",
    "    counts = numpy.bincount(column)\n",
    "    # Divide by the total column length to get a probability.\n",
    "    probabilities = counts / float(len(column))\n",
    "    \n",
    "    # Initialize the entropy to 0.\n",
    "    entropy = 0\n",
    "    # Loop through the probabilities, and add each one to the total entropy.\n",
    "    for prob in probabilities:\n",
    "        if prob > 0:\n",
    "            entropy += prob * math.log(float(prob), 2)\n",
    "    \n",
    "    return -entropy\n",
    "\n",
    "# Verify our function matches our answer from earlier.\n",
    "entropy = calc_entropy([1,1,0,0,1])\n",
    "#print(entropy)\n",
    "\n",
    "information_gain = entropy - ((.8 * calc_entropy([1,1,0,0])) + (.2 * calc_entropy([1])))\n",
    "#print(information_gain)\n",
    "income_entropy = calc_entropy(income[\"high_income\"])\n",
    "print income_entropy\n",
    "median_age = income[\"age\"].median()\n",
    "\n",
    "left_split = income[income[\"age\"] <= median_age]\n",
    "right_split = income[income[\"age\"] > median_age]\n",
    "\n",
    "age_information_gain = income_entropy - ((left_split.shape[0] / income.shape[0]) * calc_entropy(left_split[\"high_income\"]) + ((right_split.shape[0] / income.shape[0]) * calc_entropy(right_split[\"high_income\"])))\n",
    "#print age_information_gain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 3 0 2 0 1]\n"
     ]
    }
   ],
   "source": [
    "print numpy.bincount([1,1,0,0,1,3,3,5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0501271848501\n"
     ]
    }
   ],
   "source": [
    "def calc_information_gain(data, split_name, target_name):\n",
    "    \"\"\"\n",
    "    Calculate information gain given a dataset, column to split on, and target.\n",
    "    \"\"\"\n",
    "    # Calculate original entropy.\n",
    "    original_entropy = calc_entropy(data[target_name])\n",
    "    \n",
    "    # Find the median of the column we're splitting.\n",
    "    column = data[split_name]\n",
    "    median = column.median()\n",
    "    \n",
    "    # Make two subsets of the data based on the median.\n",
    "    left_split = data[column <= median]\n",
    "    right_split = data[column > median]\n",
    "    \n",
    "    # Loop through the splits, and calculate the subset entropy.\n",
    "    to_subtract = 0\n",
    "    for subset in [left_split, right_split]:\n",
    "        prob = (subset.shape[0] / float(data.shape[0])) \n",
    "        to_subtract += prob * calc_entropy(subset[target_name])\n",
    "    \n",
    "    # Return information gain.\n",
    "    return original_entropy - to_subtract\n",
    "\n",
    "# Verify that our answer is the same as in the last screen.\n",
    "print(calc_information_gain(income, \"age\", \"high_income\"))\n",
    "\n",
    "columns = [\"age\", \"workclass\", \"education_num\", \"marital_status\", \"occupation\", \"relationship\", \"race\", \"sex\", \"hours_per_week\", \"native_country\"]\n",
    "information_gains = []\n",
    "# Loop through and compute information gains.\n",
    "for col in columns:\n",
    "    information_gain = calc_information_gain(income, col, \"high_income\")\n",
    "    information_gains.append(information_gain)\n",
    "\n",
    "# Find the name of the column with the highest gain.\n",
    "highest_gain_index = information_gains.index(max(information_gains))\n",
    "highest_gain = columns[highest_gain_index]\n",
    "#print highest_gain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "marital_status\n"
     ]
    }
   ],
   "source": [
    "def find_best_column(data, target_name, columns):\n",
    "    information_gains = []\n",
    "    # Loop through and compute information gains.\n",
    "    for col in columns:\n",
    "        information_gain = calc_information_gain(data, col, \"high_income\")\n",
    "        information_gains.append(information_gain)\n",
    "\n",
    "    # Find the name of the column with the highest gain.\n",
    "    highest_gain_index = information_gains.index(max(information_gains))\n",
    "    highest_gain = columns[highest_gain_index]\n",
    "    return highest_gain\n",
    "\n",
    "income_split = find_best_column(income, \"high_income\", columns)\n",
    "print income_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# We'll use lists to store our labels for nodes (when we find them).\n",
    "# Lists can be accessed inside our recursive function, whereas integers can't.  \n",
    "# Look at the python missions on scoping for more information on this.\n",
    "def id3(data, target, columns):\n",
    "    \n",
    "    unique_targets = pandas.unique(data[target])\n",
    "\n",
    "    if len(unique_targets) == 1:\n",
    "        if 0 in unique_targets:\n",
    "            label_0s.append(0)\n",
    "        elif 1 in unique_targets:\n",
    "            label_1s.append(1)\n",
    "        return\n",
    "    \n",
    "    best_column = find_best_column(data, target, columns)\n",
    "    column_median = data[best_column].median()\n",
    "    \n",
    "    left_split = data[data[best_column] <= column_median]\n",
    "    right_split = data[data[best_column] > column_median]\n",
    "    \n",
    "    for split in [left_split, right_split]:\n",
    "        id3(split, target, columns)\n",
    "data[['marital_status', 'age','high_income']] = data[['marital_status', 'age','high_income']].astype('int32')\n",
    "\n",
    "\n",
    "id3(data, \"high_income\", [\"age\", \"marital_status\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Create a dictionary to hold the tree.  This has to be outside the function so we can access it later.\n",
    "tree = {}\n",
    "\n",
    "# This list will let us number the nodes.  It has to be a list so we can access it inside the function.\n",
    "nodes = []\n",
    "def id3(data, target, columns, tree):\n",
    "    unique_targets = pandas.unique(data[target])\n",
    "    nodes.append(len(nodes) + 1)\n",
    "    tree[\"number\"] = nodes[-1]\n",
    "\n",
    "    if len(unique_targets) == 1:\n",
    "        if 0 in unique_targets:\n",
    "            tree[\"label\"] = 0\n",
    "        elif 1 in unique_targets:\n",
    "            tree[\"label\"] = 1\n",
    "        return\n",
    "    \n",
    "    best_column = find_best_column(data, target, columns)\n",
    "    column_median = data[best_column].median()\n",
    "    \n",
    "    tree[\"column\"] = best_column\n",
    "    tree[\"median\"] = column_median\n",
    "    \n",
    "    left_split = data[data[best_column] <= column_median]\n",
    "    right_split = data[data[best_column] > column_median]\n",
    "    split_dict = [[\"left\", left_split], [\"right\", right_split]]\n",
    "    \n",
    "    for name, split in split_dict:\n",
    "        tree[name] = {}\n",
    "        id3(split, target, columns, tree[name])\n",
    "\n",
    "\n",
    "id3(data, \"high_income\", [\"age\", \"marital_status\"], tree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "age > 37.5\n",
      "age > 37.5\n",
      "    age > 25.0\n",
      "        age > 22.5\n",
      "            Leaf: Label 0\n",
      "            Leaf: Label 1\n",
      "        Leaf: Label 1\n",
      "    age > 55.0\n",
      "        age > 47.5\n",
      "            Leaf: Label 0\n",
      "            Leaf: Label 1\n",
      "        Leaf: Label 0\n"
     ]
    }
   ],
   "source": [
    "def print_with_depth(string, depth):\n",
    "    # Add space before a string.\n",
    "    prefix = \"    \" * depth\n",
    "    # Print a string, appropriately indented.\n",
    "    print(\"{0}{1}\".format(prefix, string))\n",
    "    \n",
    "    \n",
    "def print_node(tree, depth):\n",
    "    # Check for the presence of label in the tree.\n",
    "    if \"label\" in tree:\n",
    "        # If there's a label, then this is a leaf, so print it and return.\n",
    "        print_with_depth(\"Leaf: Label {0}\".format(tree[\"label\"]), depth)\n",
    "        # This is critical -- without it, you'll get infinite recursion.\n",
    "        return\n",
    "    # Print information about what the node is splitting on.\n",
    "    print_with_depth(\"{0} > {1}\".format(tree[\"column\"], tree[\"median\"]), depth)\n",
    "    \n",
    "    # Create a list of tree branches.\n",
    "    branches = [tree[\"left\"], tree[\"right\"]]\n",
    "        \n",
    "    # Insert code here to recursively call print_node on each branch.\n",
    "    # Don't forget to increment depth when you pass it in!\n",
    "\n",
    "print_node(tree, 0)\n",
    "def print_node(tree, depth):\n",
    "    if \"label\" in tree:\n",
    "        print_with_depth(\"Leaf: Label {0}\".format(tree[\"label\"]), depth)\n",
    "        return\n",
    "    print_with_depth(\"{0} > {1}\".format(tree[\"column\"], tree[\"median\"]), depth)\n",
    "    for branch in [tree[\"left\"], tree[\"right\"]]:\n",
    "        print_node(branch, depth+1)\n",
    "\n",
    "print_node(tree, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "def predict(tree, row):\n",
    "    if \"label\" in tree:\n",
    "        return tree[\"label\"]\n",
    "    \n",
    "    column = tree[\"column\"]\n",
    "    median = tree[\"median\"]\n",
    "    \n",
    "    # Insert code here to check if row[column] is less than or equal to median\n",
    "    # If it's less than or equal, return the result of predicting on the left branch of the tree\n",
    "    # If it's greater, return the result of predicting on the right branch of the tree\n",
    "    # Remember to use the return statement to return the result!\n",
    "\n",
    "# Print the prediction for the first row in our data.\n",
    "print(predict(tree, data.iloc[0]))\n",
    "def predict(tree, row):\n",
    "    if \"label\" in tree:\n",
    "        return tree[\"label\"]\n",
    "    \n",
    "    column = tree[\"column\"]\n",
    "    median = tree[\"median\"]\n",
    "    if row[column] <= median:\n",
    "        return predict(tree[\"left\"], row)\n",
    "    else:\n",
    "        return predict(tree[\"right\"], row)\n",
    "\n",
    "print(predict(tree, data.iloc[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0    0\n",
      "1    0\n",
      "2    0\n",
      "3    0\n",
      "4    1\n",
      "5    0\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "new_data = pandas.DataFrame([\n",
    "    [40,0],\n",
    "    [20,2],\n",
    "    [80,1],\n",
    "    [15,1],\n",
    "    [27,2],\n",
    "    [38,1]\n",
    "    ])\n",
    "# Assign column names to the data.\n",
    "new_data.columns = [\"age\", \"marital_status\"]\n",
    "\n",
    "def batch_predict(tree, df):\n",
    "    # Insert your code here.\n",
    "    pass\n",
    "\n",
    "predictions = batch_predict(tree, new_data)\n",
    "def batch_predict(tree, df):\n",
    "    return df.apply(lambda x: predict(tree, x), axis=1)\n",
    "\n",
    "predictions = batch_predict(tree, new_data)\n",
    "print predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
